import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision import transforms
from torchvision.transforms import ToTensor
import torchvision.transforms as tt
import numpy as np
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn import metrics

import os

# 导入自己创建的python文件
import sys
sys.path.append("..") # Adds higher directory to python modules path.
from frame.DataProcess import *
from frame.TrainUtil import *
from frame.LIRAAttack import *
from frame.AttackUtil import *
from frame.ShadowAttack import *
from frame.ThresholdAttack import *
from frame.LabelAttack import *


LEARNING_RATE = 1e-3
BATCH_SIZE = 64
MODEL = 'CNN'
EPOCHS = 50
DATA_NAME = 'CIFAR10' 
weight_dir = os.path.join('..', 'weights_for_exp', DATA_NAME)
num_shadowsets = 100
seed = 0
prop_keep = 0.5

model_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.507, 0.487, 0.441], std=[0.267, 0.256, 0.276])
    ])
attack_transform = transforms.Compose([])
device = "cuda" if torch.cuda.is_available() else "cpu"

# 影子模型攻击相关参数
sha_models = [1,2,3] #[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30]
tar_model = 0
attack_class = False #是否针对每个类别分别攻击
attack_lr = 5e-4

X_data, Y_data, train_keep = load_CIFAR10_keep(num_shadowsets, prop_keep, seed)

train_data = CustomDataset(X_data, Y_data, model_transform)
train_dataloader = DataLoader(train_data, batch_size=64, shuffle=False)

batch_size = BATCH_SIZE
model = MODEL
epochs = EPOCHS
data_name = DATA_NAME 

loss_data_all = np.load('CIFAR10_loss.npy')
score_all = np.load('CIFAR10_score.npy')
conf_data_all = np.load('CIFAR10_conf.npy')

pri_risk_all = get_risk_score(loss_data_all, train_keep)

pri_risk_rank = np.argsort(pri_risk_all)
pri_risk_rank = np.flip(pri_risk_rank)

logits_data_all = np.load('CIFAR10_logits.npy')

# 按照k个模型进行拼接
k = 10
for i in range(k):
    if i == 0:
        combine_features = logits_data_all[i]
    else:
        combine_features = np.concatenate((combine_features, logits_data_all[i]),axis=1)

# 数据量太大，不能保存所有的余弦相似度，只能需要时计算
alpha_list = [0.05, 0.1, 0.12, 0.15, 0.2, 0.3]
n_num_list = []
for i in range(combine_features.shape[0]):
# for i in range(10000):
    n_count = [0 for _ in alpha_list]
    if i%50 == 0:
        print(f"compute to: {i}")
    for j in range(combine_features.shape[0]):
        # 余弦距离的计算
        vec1 = combine_features[i]
        vec2 = combine_features[j]        
        cos_sim = vec1.dot(vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
        cos_dis = 0.5 - 0.5 * cos_sim
        for m in range(len(alpha_list)):
            if (cos_dis < alpha_list[m]):
                n_count[m] += 1
    n_num_list.append(n_count)


neigh_data_all = np.array(n_num_list)
np.save('CIFAR10_neigh.npy', neigh_data_all)